"""
@Time    : 2018/11/7 17:41
@Author  : CcH
"""
from random import shuffle
from DataUtils.Common import pad,unk
import numpy as np
import json
class DataLoader:
    def __init__(self,max_sentence_length,min_words):
        print('---read data---')
        with open(r'./DataUtils/corpus.conv','r',encoding='utf-8')as f:
            data=f.readlines()
        shuffle(data)
        self.min_words=min_words
        self.max_sentence_length = max_sentence_length
        self.train_content, self.test_content, self.train_label, self.test_label=self.get_data(data)
        self.char_2_idx,self.idx_2_char=self.build_vocab(data)
        self.train_x=self.sentence_2_idx(self.train_content)
        self.test_x=self.sentence_2_idx(self.test_content)
        self.train_y=self.label_2_numpy(self.train_label)
        self.test_y=self.label_2_numpy(self.test_label)

    def build_vocab(self,data):
        vocab={}
        for dt in data:
            for word in dt.strip().split(',')[0].split():
                if word in vocab:
                    vocab[word]+=1
                else:
                    vocab[word]=1
        char_2_idx={
            pad:0,
            unk:1
        }
        ix=0
        for k,v in vocab.items():
            if v >= self.min_words:
                char_2_idx[k]=ix+2
                ix+=1
        idx_2_char={}
        for k,v in char_2_idx.items():
            idx_2_char[v]=k
        with open('word_2_idx.json','w',encoding='utf-8')as fw:
            json.dump(char_2_idx, fw, ensure_ascii=False)
        return char_2_idx,idx_2_char


    @staticmethod
    def get_data(data):
        i=0
        train_data = []
        train_label = []
        test_data = []
        test_label = []
        for dt in data:
            content_lable=dt.strip().split(',')
            if i<10000:
                try:
                    test_data.append(content_lable[0])
                    test_label.append(content_lable[1])
                except:
                    print(i)
            else:
                try:
                    train_data.append(content_lable[0])
                    train_label.append(content_lable[1])
                except:
                    print(i)
            i+=1
        return train_data,test_data,train_label,test_label

    def sentence_2_idx(self,content):
        text_idx=[]
        for ct in content:
            sents= ct.split()
            sent_idx = []
            if len(sents) >= self.max_sentence_length:
                for word in sents[:self.max_sentence_length]:
                    if word in self.char_2_idx:
                        sent_idx.append(self.char_2_idx[word])
                    else:
                        sent_idx.append(1)
            else:
                for word in sents:
                    if word in self.char_2_idx:
                        sent_idx.append(self.char_2_idx[word])
                    else:
                        sent_idx.append(1)
                sent_idx=sent_idx+[0]*(self.max_sentence_length-len(sents))
                # sent_idx=[self.char_2_idx[word] for word in sents]+[0]*(self.max_sentence_length-len(sents))
            text_idx.append(sent_idx)
        return text_idx
    def label_2_numpy(self,labels):
        label_list=[]
        for i in range(len(labels)):
            if labels[i]=='negative':
                label_list.append(0)

            elif labels[i]=='neutral':
                label_list.append(1)

            else:
                label_list.append(2)   #0 is negative ,1 is positive

        return label_list


if __name__ == '__main__':
    data_loader=DataLoader(max_sentence_length=10,min_words=6)
    # print(data_loader.test_x)
    length=len(data_loader.char_2_idx)
    print(length)
    print(data_loader.train_label[:10])
    print(data_loader.train_y[:10])

